import dashscope
from http import HTTPStatus


class BaseAgent:
    def __init__(self, model, **kwargs):
        """
        初始化基础代理类。

        参数:
        - model: str，指定要使用的模型。
        - **kwargs: 额外的关键字参数。
        """
        self.model = model  # 保存指定的模型
        self.api_key = "sk-rnqpCFMhgP"  # 设置DashScope的API密钥

    def _call_llm(self, messages, **kwargs):
        """
        调用语言模型并返回生成的文本。

        参数:
        - messages: List[dict]，传递给模型的消息列表，每个消息是一个字典。
        - **kwargs: 额外的关键字参数，可以包括top_p等。

        返回:
        - str，模型生成的文本或者错误信息。
        """
        # 设置top_p参数，默认值为0.8
        top_p = kwargs.get('top_p', 0.8)
        # 设置DashScope的API密钥
        dashscope.api_key = self.api_key
        # 调用DashScope的生成模型接口
        response = dashscope.Generation.call(
            self.model,  # 指定使用的模型
            messages=messages,  # 传递给模型的消息
            result_format='message',  # 指定结果格式为消息
            stream=False,  # 不使用流式传输
            top_p=top_p,  # 设置top_p参数
        )

        # 检查响应状态码是否为200，表示请求成功
        if response.status_code == HTTPStatus.OK:
            # 记录最后一次调用的token信息
            # self.stat_last_call_token_info_no_stream(response)
            # 返回模型生成的内容
            return response.output.choices[0].message.content
        else:
            # 构造错误信息
            err = 'Error code: %s, error message: %s' % (
                response.code,
                response.message,
            )
            # 返回错误信息
            return err

    def stat_last_call_token_info_no_stream(self, response):
        """
        统计并保存最后一次调用的令牌使用信息。

        该方法从API响应中提取令牌使用信息，包括输入令牌、输出令牌和总令牌数，
        并保存这些信息以便后续使用。如果响应中不包含usage信息，它会打印提示。

        参数:
        response (object): API响应对象，应包含usage属性，该属性描述了令牌的使用情况。

        返回:
        object: 返回经过处理的响应对象，同时更新了类内部的last_call_usage_info属性。
        """
        try:
            # 检查response对象中是否包含usage信息
            if response.usage is not None:
                # 根据响应对象中的usage信息计算总令牌数
                if not response.usage.get('total_tokens'):
                    total_tokens = response.usage.input_tokens + response.usage.usage.output_tokens
                else:
                    total_tokens = response.usage.total_tokens

                # 保存令牌使用信息到类的属性中
                self.last_call_usage_info = {
                    'prompt_tokens': response.usage.input_tokens,
                    'completion_tokens': response.usage.output_tokens,
                    'total_tokens': total_tokens
                }
            else:
                # 如果response中不包含usage信息，打印提示
                print('No usage info in response')
        except AttributeError:
            # 捕获AttributeError，这可能发生在response没有usage属性时
            print('No usage info in response')
        # 返回处理后的响应对象
        return response